'''
Created on Jul 27, 2012

@author: pedro
'''
import matplotlib.pyplot as plt
import numpy as np
import Utils

def save_to_heatmap(data,filename):
    data = np.array(data)
    
    plt.figure(1)
    plt.hot()
    plt.pcolormesh(data)
    plt.colorbar()
    plt.savefig(filename)
    plt.close()

    
def avg_output_bar_graph(bminus, pminus, bplus, pplus, bb_pstat, pp_pstat, filename=None):
    
    b_minus_mean = Utils.calc_mean(bminus)
    b_plus_mean = Utils.calc_mean(bplus)
    p_minus_mean = Utils.calc_mean(pminus)
    p_plus_mean = Utils.calc_mean(pplus)
    
    b_minus_std = Utils.stdv(bminus)
    b_plus_std  = Utils.stdv(bplus)
    p_minus_std = Utils.stdv(pminus)
    p_plus_std  = Utils.stdv(pplus)
    
    xaxis_width = 1
    ytick_width = 5
    
    max_val = (2*ytick_width) + max([p_minus_mean  + p_minus_std, 
                                 b_minus_mean  + b_minus_std,
                                 p_plus_mean   + p_plus_std,
                                 b_plus_mean   + b_plus_std])
    
    min_val = (2*-ytick_width) + max([p_minus_mean - p_minus_std, 
                                  b_minus_mean  - b_minus_std, 
                                  p_plus_mean   - p_plus_std,
                                  b_plus_mean   - b_plus_std])
    
    max_val = int(max_val)
    min_val = int(min_val)
    
    fig2 = plt.figure(2)
    
    b_m = plt.bar(0, b_minus_mean, xaxis_width, color='r', yerr=b_minus_std)
    b_p = plt.bar(1, b_plus_mean,  xaxis_width, color='y', yerr=b_plus_std)
    p_m = plt.bar(3, p_minus_mean, xaxis_width, color='g', yerr=p_minus_std)
    p_p = plt.bar(4, p_plus_mean,  xaxis_width, color='b', yerr=p_plus_std)

    plt.ylabel('Avg. Neural Output')
    plt.title('Avg. Neural Output per Trained (+) v.s. Untrained (-) HDRs:\nBefore(B) v.s. After(P) Extra Training')
    plt.yticks(np.arange(min_val,max_val,ytick_width))
    plt.legend( (b_m[0], b_p[0], p_m[0], p_p[0]), ('B-', 'B+', 'P-', 'P+') )
    plt.axis([0,ytick_width,min_val,max_val])
    plt.xticks(np.arange(4),("","","",""))
    
    props = {'connectionstyle':'bar','arrowstyle':'-',\
                 'shrinkA':20,'shrinkB':20,'lw':2}
    
    if bb_pstat < 0.10:
#        plt.annotate('*', xy=(1.5, b_plus_mean+ytick_width), xytext=(1.5, b_plus_mean+ytick_width+1),
#            arrowprops=props)
        plt.annotate("* p < 0.10", xy=(1, b_plus_mean+4), zorder=10)
        plt.annotate('', xy=(0, b_plus_mean), xytext=(2, b_plus_mean), arrowprops=props)
        
    if pp_pstat < 0.10:
#        plt.annotate('*', xy=(3.5, p_minus_std+ytick_width), xytext=(3.5, p_minus_std+ytick_width+1),
#            arrowprops=props)
        plt.annotate("* p < 0.10", xy=(4, b_plus_mean+4), zorder=10)
        plt.annotate('', xy=(3, b_plus_mean), xytext=(5, b_plus_mean), arrowprops=props)
        
    
    if filename == None:
        plt.show()
    else:
        plt.savefig(filename)